#include "vk_texture.h"
#include "stb_image.h"
#include <iostream>
#include "vk_memory.h"


namespace vkrender
{
    VulkanTexture::VulkanTexture(const std::shared_ptr<VulkanDeviceContext>& context,
        const std::string& filepath)
    {
        this->context = context;
        loadTexture(filepath);
        createImageView();
    }

    VulkanTexture::~VulkanTexture()
    {
        vkDestroyImageView(context->device, imageView, nullptr);
        vkDestroyImage(context->device, image, nullptr);
        vkFreeMemory(context->device, memory, nullptr);
    }

    bool VulkanTexture::loadTexture(const std::string& filepath)
    {
        // load image
        int width, height, channels;
        stbi_uc* pixels = stbi_load(filepath.c_str(), &width, &height, &channels, STBI_rgb_alpha);
        if (!pixels)
        {
            std::cout << "failed to load texture\n";
            return false;
        }
        VkDeviceSize image_size = width * height * 4;
        
        // use staging buffer for better performance in rendering
        VkBuffer stagingBuffer;
        VkDeviceMemory stagingMemory;
        this->context->createBuffer(
            // context->device, context->physicalDevice, 
            image_size,
            VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, 
            stagingBuffer, stagingMemory);
        // copy pixel data from host to staging buffer
        void* data;
        vkMapMemory(context->device, stagingMemory, 0, image_size, 0, &data);
        memcpy(data, pixels, static_cast<size_t>(image_size));
        vkUnmapMemory(context->device, stagingMemory);
        stbi_image_free(pixels); 

        // create image and memory
        VkImageCreateInfo imageInfo = {};
        imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
        imageInfo.imageType = VK_IMAGE_TYPE_2D;
        imageInfo.extent.width = width;
        imageInfo.extent.height = height;
        imageInfo.extent.depth = 1;
        imageInfo.mipLevels = 1;
        imageInfo.arrayLayers = 1;
        imageInfo.format = VK_FORMAT_R8G8B8A8_SRGB;
        imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
        imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
        imageInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
        imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
        imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

        createImageWithInfo(context->device, context->physicalDevice, imageInfo,
            VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, this->image, this->memory);

        context->transitionImageLayout(image, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
        context->copyBufferToImage(stagingBuffer, image, width, height);
        context->transitionImageLayout(image, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

        vkDestroyBuffer(context->device, stagingBuffer, nullptr);
        vkFreeMemory(context->device, stagingMemory, nullptr);

        return true;
    }

    void VulkanTexture::createImageView()
    {
        imageView = context->createImageView(image, VK_FORMAT_R8G8B8A8_SRGB);
    }
}
